/*
 * Copyright (C) 2020 original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 * in compliance with the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the License
 * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 * or implied. See the License for the specific language governing permissions and limitations under
 * the License.
 */

package vip.justlive.easyproxy.server.handler;

import java.nio.ByteBuffer;
import java.util.EnumMap;
import java.util.function.BiConsumer;
import lombok.extern.slf4j.Slf4j;
import vip.justlive.easyproxy.core.context.Metrics;
import vip.justlive.easyproxy.core.context.RemoteContext;
import vip.justlive.easyproxy.core.context.ServerProxyContext;
import vip.justlive.easyproxy.core.protocol.AbstractProxyDataHandler;
import vip.justlive.easyproxy.core.protocol.ChannelData;
import vip.justlive.easyproxy.core.protocol.ProxyData;
import vip.justlive.easyproxy.core.protocol.ProxyDataType;
import vip.justlive.oxygen.core.net.aio.core.AioListener;
import vip.justlive.oxygen.core.net.aio.core.ChannelContext;
import vip.justlive.oxygen.core.util.Strings;
import vip.justlive.oxygen.ioc.annotation.Bean;

/**
 * 服务端处理单元
 *
 * @author wubo
 */
@Bean
@Slf4j
public class ServerHandler extends AbstractProxyDataHandler implements AioListener {

  private final EnumMap<ProxyDataType, BiConsumer<ProxyData, ChannelContext>> consumers = new EnumMap<>(
      ProxyDataType.class);

  public ServerHandler() {
    consumers.put(ProxyDataType.AUTH, this::handleAuth);
    consumers.put(ProxyDataType.TRANSFER, this::handleTransfer);
    consumers.put(ProxyDataType.CLOSE_CHANNEL, this::handleCloseChannel);
  }

  @Override
  public void doHandle(ProxyData data, ChannelContext channelContext) {
    BiConsumer<ProxyData, ChannelContext> consumer = consumers.get(data.getType());
    if (consumer == null) {
      log.warn("no consumer found for {}", data.getType());
      return;
    }
    consumer.accept(data, channelContext);
  }

  @Override
  public void onConnected(ChannelContext channelContext) {
    Metrics.metrics(channelContext.getServerAddress().getPort()).getServerChannels()
        .set(channelContext.getGroupContext().getChannels().size());
  }

  @Override
  public void onClosed(ChannelContext channelContext) {
    Metrics.metrics(channelContext.getServerAddress().getPort()).getServerChannels()
        .set(channelContext.getGroupContext().getChannels().size());
    RemoteContext.remove(channelContext);
  }

  @Override
  public ByteBuffer encode(Object data, ChannelContext channelContext) {
    ByteBuffer buffer = super.encode(data, channelContext);
    Metrics.metrics(channelContext.getServerAddress().getPort()).getWriteBytes()
        .addAndGet(buffer.limit());
    return buffer;
  }

  @Override
  public Object decode(ByteBuffer buffer, int readableSize, ChannelContext channelContext) {
    int position = buffer.position();
    Object obj = super.decode(buffer, readableSize, channelContext);
    if (obj != null) {
      Metrics.metrics(channelContext.getServerAddress().getPort()).getReadBytes()
          .addAndGet(buffer.position() - position);
    }
    return obj;
  }

  private void handleAuth(ProxyData data, ChannelContext channelContext) {
    String authKey = data.bodyAsString();
    if (!Strings.hasText(authKey)) {
      channelContext.close();
      return;
    }
    RemoteContext.register(authKey, channelContext);
  }

  private void handleTransfer(ProxyData data, ChannelContext channelContext) {
    ServerProxyContext proxy = proxy(data, channelContext);
    if (proxy == null) {
      return;
    }

    ChannelData cd = data.deserialize();
    ChannelContext channel = proxy.getServer().getGroupContext().getChannels().get(cd.getId());
    if (channel == null) {
      log.error("channel not found {} {} {}/bytes", data.getId(), cd.getId(), cd.getData().length);
      return;
    }
    channel.write(cd.getData());
  }

  private void handleCloseChannel(ProxyData data, ChannelContext channelContext) {
    ServerProxyContext proxy = proxy(data, channelContext);
    if (proxy != null) {
      ChannelContext channel = proxy.getServer().getGroupContext().getChannels()
          .get(data.deserialize().getId());
      if (channel != null) {
        channel.close();
      }
    }
  }

  private ServerProxyContext proxy(ProxyData data, ChannelContext channelContext) {
    RemoteContext remote = RemoteContext.lookup(channelContext);
    if (remote == null) {
      log.error("remote not found {}", channelContext);
      channelContext.close();
      return null;
    }

    ServerProxyContext proxy = remote.find(data.getId());
    if (proxy == null) {
      log.error("proxy not found {}", data.getId());
    }
    return proxy;
  }

}
